import os
import logging
import time
import config_extraction
from kubernetes import client
from pyfiglet import Figlet


def update_taint_on_node(v1, node_name, taint_key="high-memory", taint_value="true", taint_effect="NoSchedule",
                         action="add"):
    """
    更新节点上的污点。
    :param v1: CoreV1Api 实例
    :param node_name: 节点名称
    :param taint_key: 污点键
    :param taint_value: 污点值
    :param taint_effect: 污点效果
    :param action: 动作，可以是 'add' 或 'remove'
    """
    try:
        # 获取节点对象
        node = v1.read_node(node_name)
        existing_taints = node.spec.taints if node.spec.taints else []

        if action == "add":
            # 检查是否存在相同的污点
            for taint in existing_taints:
                if taint.key == taint_key and taint.value == taint_value and taint.effect == taint_effect:
                    logging.info(f"Taint already exists on node {node_name}: {taint_key}={taint_value}: {taint_effect}")
                    return

            # 添加新的污点
            new_taint = client.V1Taint(key=taint_key, value=taint_value, effect=taint_effect)
            body = {
                "spec": {
                    "taints": existing_taints + [new_taint]
                }
            }
            v1.patch_node(node_name, body)
            logging.info(f"Taint added to node {node_name}: {taint_key}={taint_value}: {taint_effect}")

        elif action == "remove":
            # 移除指定的污点
            updated_taints = [
                taint for taint in existing_taints
                if not (taint.key == taint_key and taint.value == taint_value and taint.effect == taint_effect)
            ]

            if len(updated_taints) != len(existing_taints):
                body = {
                    "spec": {
                        "taints": updated_taints
                    }
                }
                v1.patch_node(node_name, body)
                logging.info(f"Taint removed from node {node_name}: {taint_key}={taint_value}: {taint_effect}")
            else:
                logging.debug(f"Taint does not exist on node {node_name}: {taint_key}={taint_value}: {taint_effect}")

    except Exception as e:
        logging.error(f"Failed to update taint on node {node_name}: {e}")


def restart_resource(apps_v1, namespace, resource_type, resource_name):
    """
    重启资源
    :param apps_v1:
    :param namespace:
    :param resource_type: 资源类型(Deployment、StatefulSet、DaemonSet、ReplicaSet)
    :param resource_name: 资源名称
    :return:
    """
    try:
        body = {
            "spec": {
                "template": {
                    "metadata": {
                        "annotations": {
                            "kubectl.kubernetes.io/restartedAt": time.strftime(
                                "%Y-%m-%dT%H:%M:%SZ", time.gmtime()
                            )
                        }
                    }
                }
            }
        }
        if resource_type == 'deployment':
            apps_v1.patch_namespaced_deployment(name=resource_name, namespace=namespace, body=body)
        elif resource_type == 'statefulset':
            apps_v1.patch_namespaced_stateful_set(name=resource_name, namespace=namespace, body=body)
        elif resource_type == 'daemonset':
            apps_v1.patch_namespaced_daemon_set(name=resource_name, namespace=namespace, body=body)
        elif resource_type == 'replicaset':
            apps_v1.patch_namespaced_replica_set(name=resource_name, namespace=namespace, body=body)
        else:
            logging.warning(f"Unsupported resource type: {resource_type}")
            return
        logging.info(f"Restarted {resource_type} {resource_name} in namespace {namespace}")
    except client.exceptions.ApiException as e:
        logging.error(f"Failed to restart {resource_type} {resource_name} in namespace {namespace}: {e}")


def is_whitelisted(namespace, controller_kind, name, whitelist):
    """
    检查命名空间、控制器类型和名称是否在白名单中。
    :param namespace: 命名空间
    :param controller_kind: 控制器类型（如 deployment, StatefulSet, DaemonSet）
    :param name: 控制器名称
    :param whitelist: 白名单字典
    :return: 是否在白名单中
    """
    if not whitelist:
        logging.warning("Whitelist is empty or not configured.")
        return False

    namespace = namespace.lower()
    controller_kind = controller_kind.lower()
    name = name.lower()

    # Convert inputs to lowercase for case-insensitive comparison
    if namespace in whitelist:
        if isinstance(whitelist[namespace], dict):
            if controller_kind in whitelist[namespace]:
                if (whitelist[namespace][controller_kind] == '*' or
                        whitelist[namespace][controller_kind].lower() == name):
                    return True
        elif whitelist[namespace] == '*':
            return True

    # 检查全局通配符
    if '*' in whitelist:
        if whitelist['*'] == '*':
            return True

    return False


def handle_high_memory_node(v1, custom_objects_api, apps_v1, host, cfg):
    """
    处理高内存节点。
    :param v1: CoreV1Api 实例
    :param custom_objects_api: CustomObjectsApi 实例
    :param apps_v1: AppsV1Api 实例
    :param host: 主机地址
    :param cfg: 配置字典
    """
    # 获取实际的 URL 从配置文件
    actual_host = cfg['aliases'].get(host)

    if not actual_host:
        logging.error(f"No alias found for host '{host}'.")
        return

    if not actual_host.startswith(('http://', 'https://')):
        logging.error(f"Invalid URL '{actual_host}' for host '{host}'. No scheme supplied.")
        return

    metrics = config_extraction.get_node_metrics(actual_host)
    if not metrics:
        return

    mem_usage = config_extraction.parse_metrics(metrics)
    logging.info(f"Node {host} memory usage: {mem_usage}%")

    # 提取别名中的 IP 地址
    ip_address = actual_host.split(':')[1].strip('//')

    # 查找节点名称
    node_name = None
    for node in v1.list_node().items:
        internal_ip = None

        for address in node.status.addresses:
            if address.type == 'InternalIP':
                internal_ip = address.address
                break

        logging.debug(f"Checking node {node.metadata.name}: InternalIP={internal_ip}")

        if internal_ip and internal_ip == ip_address:
            node_name = node.metadata.name
            break

    if not node_name:
        logging.error(f"Node name not found for IP: {ip_address}")
        return

    if mem_usage <= cfg['mem_threshold']:
        logging.debug(
            f"Memory usage {mem_usage}% is below threshold {cfg['mem_threshold']}%. Removing taint if present.")
        # 移除污点
        update_taint_on_node(v1, node_name, action="remove")
        return

    logging.info(f"Node {node_name} has high memory usage ({mem_usage}%).")

    # 添加污点
    update_taint_on_node(v1, node_name)

    # 检查所有节点是否均超过阈值
    all_nodes_exceed_threshold = True
    for node in v1.list_node().items:
        internal_ip = None
        for address in node.status.addresses:
            if address.type == 'InternalIP':
                internal_ip = address.address
                break

        if internal_ip:
            node_actual_host = cfg['aliases'].get(internal_ip)
            if node_actual_host:
                node_metrics = config_extraction.get_node_metrics(node_actual_host)
                if node_metrics:
                    node_mem_usage = config_extraction.parse_metrics(node_metrics)
                    if node_mem_usage <= cfg['mem_threshold']:
                        all_nodes_exceed_threshold = False
                        break

    if all_nodes_exceed_threshold:
        logging.info("All nodes exceed the memory threshold. Taking no further action.")
        return

    # 获取节点上的 Pods
    pods = v1.list_pod_for_all_namespaces(field_selector=f'spec.nodeName={node_name}').items
    if not pods:
        logging.info(f"No pods found on node {node_name}.")
        return

    # 过滤白名单
    filtered_pods = []
    for pod in pods:
        owner_refs = pod.metadata.owner_references
        if owner_refs:
            for ref in owner_refs:
                controller_kind = ref.kind.lower()
                controller_name = ref.name.lower()
                if not is_whitelisted(pod.metadata.namespace, controller_kind, controller_name,
                                      cfg.get('whiteList', {})):
                    filtered_pods.append((pod.metadata.namespace, controller_kind, ref.name, pod.metadata.name))
                    logging.info(
                        f"Filtered Pod: {pod.metadata.namespace}.{controller_kind}|{ref.name}|{pod.metadata.name}")

    if not filtered_pods:
        logging.info(f"No eligible pods found on node {node_name} after filtering.")
        return

    # 获取所有 Pod 的内存使用率
    pod_mem_usages = config_extraction.get_pod_mem_usages(custom_objects_api)
    max_mem_usage = 0
    target_pod = None

    for namespace, kind, name, pod_name in filtered_pods:
        if (namespace, pod_name) in pod_mem_usages:
            mem_usage = pod_mem_usages[(namespace, pod_name)]
            if mem_usage > max_mem_usage:
                max_mem_usage = mem_usage
                target_pod = (namespace, kind, name)

    if target_pod:
        namespace, kind, name = target_pod
        restart_resource(apps_v1, namespace, kind, name)
    else:
        logging.info(f"No suitable pod found to restart on node {node_name}.")


def init():
    """
    初始化
    :return:
    """
    # 加载配置文件
    config_path = os.environ.get('CONFIG_PATH', '/podBalancer/config.yaml')
    cfg = config_extraction.load_config(config_path)

    if not cfg:
        logging.error("Configuration could not be loaded.")
        return

    config_extraction.setup_logging(cfg)
    logging.debug(f"Loaded configuration: {cfg}")
    logging.info(
        f"The threshold is set to {cfg['mem_threshold']}%.")

    if not cfg.get('hosts') or not cfg.get('mem_threshold'):
        logging.error("Configuration missing required fields.")
        return

    # 初始化 Kubernetes 客户端
    try:
        v1, custom_objects_api, apps_v1 = config_extraction.load_kube_config()
    except Exception as e:
        logging.error(f"Failed to initialize Kubernetes API clients: {e}")
        return

    # 检查所有主机
    for host in cfg['hosts']:
        handle_high_memory_node(v1, custom_objects_api, apps_v1, host, cfg)

    logging.info("Monitoring completed.")


def generate_ascii_art(text, font="standard"):
    fig = Figlet(font=font)
    ascii_art = fig.renderText(text)
    print(ascii_art)


if __name__ == '__main__':
    generate_ascii_art("PodBalancer", font="block")
    init()
